import numpy as np
import time
from pyDOE import lhs
import matplotlib
# matplotlib.use('Agg')
import matplotlib.pyplot as plt
import pickle
import scipy.io
import random

# Setup GPU for training
import torch
import torch.nn as nn
import torch.optim as optim
import os

from pinn_particle_hard import sampling_from_msh
from PPINN.model import NeuralNetwork

os.environ['CUDA_VISIBLE_DEVICES'] = '0'  # CPU:-1; GPU0: 1; GPU1: 0;

random.seed(1234)
np.random.seed(1234)
torch.manual_seed(1234)
layers = np.array([2,30,30,30,30,30,30,30,30,1]) 
class PINN2D(nn.Module):
    def __init__(self, X_f, N_u, N_f, device):
        super(PINN2D, self).__init__()
        self.x_f = {}
        for key, value in X_f.items():
            self.x_f[key] = X_f[key].clone().detach().requires_grad_(True)

        layers[0] = X_f['interior'].shape[1]
        layers[-1] = 6

        self.net = NeuralNetwork(layers, device)
        self.optimizer_LBFGS = torch.optim.LBFGS(self.net.parameters(),
                                           max_iter=100000,
                                           max_eval=100000,
                                           history_size=50,
                                           tolerance_change=1 * np.finfo(float).eps,
                                           line_search_fn="strong_wolfe")
        self.optimizer_Adam = torch.optim.Adam(self.net.parameters(), lr=1e-5)
        self.ls = 0
        self.iter = 0
        self.rho = 1
        self.mu = 0.02
        
        # 接口
        Collo = np.array(self.x_f['interior'].detach().cpu().numpy())
        INLET = np.array(self.x_f['velocity_inlet_0'].detach().cpu().numpy())
        OUTLET = np.array(self.x_f['pressure_outlet_0'].detach().cpu().numpy())
        WALL = np.array(self.x_f['wall_0'].detach().cpu().numpy())
        self.uv_net = self.net

        # Collocation point
        self.x_c = torch.tensor(Collo[:, 0:1], dtype=torch.float32, requires_grad=True)
        self.y_c = torch.tensor(Collo[:, 1:2], dtype=torch.float32, requires_grad=True)

        self.x_INLET = torch.tensor(INLET[:, 0:1], dtype=torch.float32, requires_grad=True)
        self.y_INLET = torch.tensor(INLET[:, 1:2], dtype=torch.float32, requires_grad=True)

        self.x_OUTLET = torch.tensor(OUTLET[:, 0:1], dtype=torch.float32, requires_grad=True)
        self.y_OUTLET = torch.tensor(OUTLET[:, 1:2], dtype=torch.float32, requires_grad=True)

        self.x_WALL = torch.tensor(WALL[:, 0:1], dtype=torch.float32, requires_grad=True)
        self.y_WALL = torch.tensor(WALL[:, 1:2], dtype=torch.float32, requires_grad=True)

        
    def grad(self, numer, denom):
        return torch.autograd.grad(
            numer, denom,
            grad_outputs=torch.ones_like(numer),
            retain_graph=True,
            create_graph=True)[0]
        
    def u_hat(self, x, y):
        field = self.uv_net(torch.cat([x, y], 1))
        psi = field[:, 0:1]
        p = field[:, 1:2]
        sxx = field[:, 2:3]
        syy = field[:, 3:4]
        sxy = field[:, 4:5]
        ux = self.grad(psi, y)
        uy = -self.grad(psi, x)
        return ux, uy, p, sxx, syy, sxy

    def f_hat(self, x, y):
        ux, uy, p, sxx, syy, sxy = self.u_hat(x, y)

        ux_x = self.grad(ux, x) 
        ux_y = self.grad(ux, y) 

        uy_x = self.grad(uy, x) 
        uy_y = self.grad(uy, y) 
        
        sxx_x = self.grad(sxx, x) 
        
        sxy_x = self.grad(sxy, x) 
        sxy_y = self.grad(sxy, y)
        
        syy_y = self.grad(syy, y) 


        f1 = self.rho * (ux * ux_x + uy * ux_y) - sxx_x - sxy_y
        f2 = self.rho * (ux * uy_x + uy * uy_y) - sxy_x - syy_y
        f3 = self.mu * (ux_x + ux_x) - p - sxx
        f4 = self.mu * (uy_y + uy_y) - p - syy
        f5 = self.mu * (ux_y + uy_x) - sxy
        f6 = p + 0.5 * (sxx + syy) 

        return f1, f2, f3, f4, f5, f6

        
    def train(self):        
        self.net.train()
        for it in range(10000):
            self.optimizer_Adam.zero_grad()
            def closure():
                self.optimizer_LBFGS.zero_grad()
                f1, f2, f3, f4, f5, f6 = self.f_hat(self.x_c, self.y_c)
                bc1_ux, bc1_uy, _, _, _, _ = self.u_hat(self.x_WALL, self.y_WALL)
                bc2_ux, bc2_uy, _, _, _, _ = self.u_hat(self.x_INLET, self.y_INLET)
                _, _, bc3_p, _, _, _ = self.u_hat(self.x_OUTLET, self.y_OUTLET)
                
                loss_f = torch.mean(f1 ** 2) + torch.mean(f2 ** 2) + torch.mean(f3 ** 2) + torch.mean(
                    f4 ** 2) + torch.mean(f5 ** 2) + torch.mean(f6 ** 2)
                loss_WALL = torch.mean(bc1_ux ** 2) + torch.mean(bc1_uy ** 2)
                loss_INLET = torch.mean((bc2_ux - 1.0) ** 2) + torch.mean((bc2_uy - 0.0) ** 2)
                loss_OUTLET = torch.mean(bc3_p ** 2)

                loss = loss_f + 2 * (loss_WALL + loss_INLET + loss_OUTLET)
                self.ls = loss

                self.ls.backward()

                self.iter += 1
                if not self.iter % 100:
                    print('Epoch: {0:}, Loss: {1:6.9f}'.format(self.iter, self.ls))

                return self.ls
            
            self.optimizer_Adam.step()
        self.iter = 0
        self.optimizer_LBFGS.step(closure)

    def predict(self, x_star, y_star):
        x_star = torch.tensor(x_star, dtype=torch.float32, requires_grad=True)
        y_star = torch.tensor(y_star, dtype=torch.float32, requires_grad=True)
        u_star, v_star, p_star, _, _, _ = self.u_hat(x_star, y_star)
        return u_star.detach().cpu().numpy(), v_star.detach().cpu().numpy(), p_star.detach().cpu().numpy()

def postProcess(xmin, xmax, ymin, ymax, field_MIXED, s=2, alpha=0.5, marker='o'):
    [x_MIXED, y_MIXED, u_MIXED, v_MIXED, p_MIXED] = field_MIXED

    fig, ax = plt.subplots(nrows=3, ncols=1, figsize=(7, 4))
    fig.subplots_adjust(hspace=0.2, wspace=0.2)

    # Plot MIXED result
    cf = ax[0].scatter(x_MIXED, y_MIXED, c=u_MIXED, alpha=alpha-0.1, edgecolors='none', cmap='rainbow', marker=marker, s=int(s))
    ax[0].axis('square')
    for key, spine in ax[0].spines.items():
        if key in ['right','top','left','bottom']:
            spine.set_visible(False)
    ax[0].set_xticks([])
    ax[0].set_yticks([])
    ax[0].set_xlim([xmin, xmax])
    ax[0].set_ylim([ymin, ymax])
    ax[0].set_title(r'$u$ (m/s)')
    fig.colorbar(cf, ax=ax[0], fraction=0.046, pad=0.04)

    cf = ax[1].scatter(x_MIXED, y_MIXED, c=v_MIXED, alpha=alpha-0.1, edgecolors='none', cmap='rainbow', marker=marker, s=int(s))
    ax[1].axis('square')
    for key, spine in ax[1].spines.items():
        if key in ['right','top','left','bottom']:
            spine.set_visible(False)
    ax[1].set_xticks([])
    ax[1].set_yticks([])
    ax[1].set_xlim([xmin, xmax])
    ax[1].set_ylim([ymin, ymax])
    ax[1].set_title(r'$v$ (m/s)')
    fig.colorbar(cf, ax=ax[1], fraction=0.046, pad=0.04)

    cf = ax[2].scatter(x_MIXED, y_MIXED, c=p_MIXED, alpha=alpha, edgecolors='none', cmap='rainbow', marker=marker, s=int(s), vmin=-0.25, vmax=4.0)
    ax[2].axis('square')
    for key, spine in ax[2].spines.items():
        if key in ['right','top','left','bottom']:
            spine.set_visible(False)
    ax[2].set_xticks([])
    ax[2].set_yticks([])
    ax[2].set_xlim([xmin, xmax])
    ax[2].set_ylim([ymin, ymax])
    ax[2].set_title('Pressure (Pa)')
    fig.colorbar(cf, ax=ax[2], fraction=0.046, pad=0.04)

    plt.savefig('./uvp.png', dpi=300)
    plt.close('all')

if __name__ == "__main__":
    # 还原
    X_f_train = sampling_from_msh('data/PINN2D/mesh.npz', 500, 5000)
    for key, value in X_f_train.items():
        X_f_train[key] = torch.tensor(value, dtype=torch.float32, device='cuda:0')
    N_u=500
    N_f=5000

    with torch.device('cuda:0' if torch.cuda.is_available() else 'cpu'):
        # Train from scratch
        
        model = PINN2D(X_f_train, N_u, N_f, 'cuda:0')
        model.train()

        # Get mixed-form PINN prediction
        x_PINN = np.linspace(0, 1.0, 251)
        y_PINN = np.linspace(0, 0.2, 101)
        x_PINN, y_PINN = np.meshgrid(x_PINN, y_PINN)
        x_PINN = x_PINN.flatten()[:, None]
        y_PINN = y_PINN.flatten()[:, None]
        x_PINN = torch.tensor(x_PINN, dtype=torch.float32, device='cuda:0')
        y_PINN = torch.tensor(y_PINN, dtype=torch.float32, device='cuda:0')
        u_PINN, v_PINN, p_PINN = model.predict(x_PINN, y_PINN)
        x_PINN = x_PINN.detach().cpu().numpy()
        y_PINN = y_PINN.detach().cpu().numpy()
        field_MIXED = [x_PINN, y_PINN, u_PINN, v_PINN, p_PINN]
        postProcess(xmin=0, xmax=1, ymin=0, ymax=0.2, field_MIXED=field_MIXED, s=3, alpha=0.5)

